import os
import shutil
import copy
import re
from PyPDF2 import PdfReader, PdfWriter
chapters = []
head_title_list = []
def sub_chapter_split(pdf_reader,outline,current_depth,total_depth):
    for index,per_outline in enumerate(outline):
        if isinstance(per_outline,list):
            if current_depth == total_depth:
                # print(per_outline)
                for sub_outline in per_outline:
                    if isinstance(sub_outline,dict):
                        if (sub_outline['/Title'],pdf_reader.get_destination_page_number(sub_outline)) not in chapters:
                        # print(sub_outline)
                            chapters.append((sub_outline['/Title'],pdf_reader.get_destination_page_number(sub_outline)))
                continue
            sub_chapter_split(pdf_reader,per_outline,current_depth + 1,total_depth)
        else:
            if current_depth == 1 and isinstance(outline[index],dict):
                head_title_list.append(outline[index]['/Title'])
            try:
                if isinstance(outline[index+1],list) and isinstance(outline[index],dict):
                    if (per_outline['/Title'], pdf_reader.get_destination_page_number(per_outline)) not in chapters:
                        chapters.append((per_outline['/Title'], pdf_reader.get_destination_page_number(per_outline)))
            except:
                pass


def sort_title(head_title_list,chapters):
    chapters_summary = []
    for per_chapter in chapters:
        title = per_chapter[0]
        if title in head_title_list:
            try:
                if len(temp_list[1]) != 0:
                    chapters_summary.append(copy.deepcopy(temp_list))
            except:
                pass
            temp_list = [title, []]
        else:
            temp_list[1].append(per_chapter)

    if len(temp_list[1]) != 0:
        chapters_summary.append(copy.deepcopy(temp_list))

    return chapters_summary


def write_pdf(pdf_reader,chapters_summary_list,output_dir,total_pages):
    print(f"请检查{output_dir}下是否有需要保存的内容，即将删除开始新的转换生成。")
    print("按下回车以继续")
    input()
    shutil.rmtree(output_dir)
    os.mkdir(output_dir)
    os.chdir(output_dir)
    origin_path = os.path.abspath(os.path.curdir)
    for out_index,per_chapter_list in enumerate(chapters_summary_list):
        os.chdir(origin_path)
        folder_name = per_chapter_list[0]
        os.mkdir(folder_name)
        os.chdir(folder_name)
        for i,per_chapter in enumerate(per_chapter_list[1]):
            print("===============写出============")
            title = per_chapter[0]
            start_page = per_chapter[1]
            try:
                next_page_pn = per_chapter_list[1][i+1][1]
            except:
                try:
                    next_page_pn  = chapters_summary_list[out_index+1][1][0][1]
                    end_page = next_page_pn
                except:
                    next_page_pn = total_pages
                    end_page = next_page_pn
            else:
                if start_page <= next_page_pn:
                    end_page =  next_page_pn

            output_file_name = title.replace(" ","")+".pdf"
            if start_page == end_page:
                end_page+=1
            rstr = r"[\/\\\:\*\?\"\<\>\|]"  # '/ \ : * ? " < > |'
            output_file_name = re.sub(rstr, "_", output_file_name)  # 替换为下划线
            writer = PdfWriter()
            for page_num in range(start_page, end_page):
                writer.add_page(pdf_reader.pages[page_num])
            print(f"正在筹备：章节-{folder_name}  文件-{output_file_name} 开始页面-{start_page} 结束界面-{end_page}")
            with open(output_file_name, "wb") as output_pdf:
                writer.write(output_pdf)
    print("===============写出已完成========")


def split_pdf_by_chapters(input_pdf_path, output_folder,depth_total = 2):#main
    # 确保输出文件夹存在
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # 读取 PDF 文档
    reader = PdfReader(input_pdf_path)
    outlines = reader.outline  # 获取书签信息
    total_pages = len(reader.pages)

    print(f"正在解析书签，共 {total_pages} 页...")

    # 解析书签


    sub_chapter_split(reader,outlines,1,depth_total)
    print(chapters)
    chapters_summary_result = sort_title(head_title_list,chapters)
    # print(chapters_summary_result)
    write_pdf(reader,chapters_summary_result,output_dir,total_pages)
    print("按章节拆分完成！")

# 使用示例
input_pdf = r"E:\AI\mineru\弹性云主机.pdf" # 替换为你的 PDF 文件路径
output_dir = r"E:\AI\mineru\output"  # 输出文件夹路径，请确保已提前创建
split_pdf_by_chapters(input_pdf, output_dir,depth_total=1)
